/*
* main.c
*
* Copyright (c) 2006 Michael Hobbs
*
* This file is part of iTorrent.
*
* iTorrent is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation; either version 2 of the License, or
* (at your option) any later version.
*
* iTorrent is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with iTorrent; if not, write to the Free Software
* Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/syslimits.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <time.h>
#include <errno.h>

#include <mach-o/dyld.h>
#include <Security/Authorization.h>
#include <Security/AuthorizationTags.h>

/*
The caller sends us 3 arguments on the command line, a listen port, a connect
port, and a ping port. The program will then listen on the listen port and
forward all connections and data to the connect port. While it does this, it
maintains a connection to the ping port and periodically sends data through the
connection. If the ping connection goes down, it assumes the parent process has
died and will terminate the program.

This program is necessary so that it can have a root setuid, which will allow it
to listen to privileged ports. It has a routine that checks if the setuid bit is
set correctly and corrects it if it is not.
*/

#define ROOT 0
#define REPAIR "--self-repair"
#define PING_SECS 2
#define NUM_SOCKETS 32
#define PING_IDX 0
#define LISTEN_IDX 1

#define OTHER(i) ((i % 2) ? (i - 1) : (i + 1))

int fds[NUM_SOCKETS];
int nfds = 0;
char buf[1024 * 64];
int t = 1;

int e_errno(int r, int line, char * file) {
	if (r < 0) {
		fprintf(stderr, "Error %d on line %d in %s.\n", errno, line, file);
		exit(errno);
	}
	return r;
}

#define E(R) e_errno(R, __LINE__, __FILE__)

void e_osstatus(OSStatus s, int line, char * file) {
	if (s != errAuthorizationSuccess) {
		fprintf(stderr, "Error %d on line %d in %s.\n", s, line, file);
		exit(s);
	}
}

#define EX(S) e_osstatus(S, __LINE__, __FILE__)

void get_ports(int argc, char ** argv, int * p1, int * p2, int * p3) {
	char * endptr;

	if (argc == 2) {
		if (strcmp(argv[1], REPAIR) == 0) {
			return;
		}
		exit(EXIT_FAILURE);
	}
	if (argc != 4) {
		exit(EXIT_FAILURE);
	}
	*p1 = strtol(argv[1], &endptr, 10);
	if (*p1 == 0 || *endptr != '\0') {
		exit(EXIT_FAILURE);
	}
	*p2 = strtol(argv[2], &endptr, 10);
	if (*p2 == 0 || *endptr != '\0') {
		exit(EXIT_FAILURE);
	}
	*p3 = strtol(argv[3], &endptr, 10);
	if (*p3 == 0 || *endptr != '\0') {
		exit(EXIT_FAILURE);
	}
}

void set_port(struct sockaddr_in * sa, int p) {
	memset(sa, 0, sizeof(*sa));
	sa->sin_family = AF_INET;
	sa->sin_port = htons(p);
	sa->sin_addr.s_addr = htonl(INADDR_LOOPBACK);
}

void do_accept(int l, int p) {
	struct sockaddr_in sa;
	socklen_t size = sizeof(sa);
	int a = E(accept(l, (struct sockaddr *) &sa, &size));
	int c = E(socket(AF_INET, SOCK_STREAM, 0));

	if (nfds >= NUM_SOCKETS - 1) {
		close(a);
		close(c);
		return;
	}
	set_port(&sa, p);
	E(connect(c, (struct sockaddr *) &sa, sizeof(sa)));
	E(setsockopt(a, IPPROTO_TCP, TCP_NODELAY, &t, sizeof(t)));
	E(setsockopt(c, IPPROTO_TCP, TCP_NODELAY, &t, sizeof(t)));
	fds[nfds] = a;
	fds[nfds+1] = c;
	nfds += 2;
}

int do_close(int i) {
	int in = fds[i];
	int out = fds[OTHER(i)];
	int low = (i % 2) ? (i - 1) : i;

	close(in);
	close(out);
	memmove(&fds[low], &fds[low + 2], (nfds - low - 2) * sizeof(int));
	nfds -= 2;
	return 1;
}

int do_recv(int i) {
	int in = fds[i];
	int out = fds[OTHER(i)];
	int size = E(recv(in, buf, sizeof(buf), 0));

	if (!size) {
		return do_close(i);
	}
	E(send(out, buf, size, 0));
	return 0;
}

void check_setuid(int argc, char ** argv) {
	char path[PATH_MAX];
	uint32_t size = PATH_MAX;

	E(_NSGetExecutablePath(path, &size));
	if (argc == 2 && strcmp(argv[1], REPAIR) == 0) {
		E(chown(path, ROOT, ROOT));
		E(chmod(path, S_IRUSR | S_IXUSR | S_IRGRP | S_IXGRP | S_IROTH | S_IXOTH | S_ISUID));
		exit(EXIT_SUCCESS);
	}
	if (geteuid() == ROOT) {
		return;
	} else {
		AuthorizationRef ref;
		AuthorizationFlags flags;
		AuthorizationItem item;
		AuthorizationRights rights;
		char * args[] = {REPAIR, NULL};
		int status;

		item.name = kAuthorizationRightExecute;
		item.flags = 0;
		item.value = path;
		item.valueLength = strlen(path);
		rights.count = 1;
		rights.items = &item;
		flags = kAuthorizationFlagDefaults | kAuthorizationFlagExtendRights | kAuthorizationFlagInteractionAllowed;
		EX(AuthorizationCreate(&rights, kAuthorizationEmptyEnvironment, flags, &ref));
		EX(AuthorizationExecuteWithPrivileges(ref, path, kAuthorizationFlagDefaults, args, NULL));
		EX(AuthorizationFree(ref, kAuthorizationFlagDefaults));
		E(wait(&status));
		if (WIFEXITED(status)) {
			if (WEXITSTATUS(status) != EXIT_SUCCESS) {
				exit(WEXITSTATUS(status));
			}
			E(execv(path, argv));
		}
		exit(128 + WTERMSIG(status));
	}
}

int main(int argc, char ** argv) {
	int p1;
	int p2;
	int p3;
	struct sockaddr_in sa;
	int c = E(socket(AF_INET, SOCK_STREAM, 0));
	int l = E(socket(AF_INET, SOCK_STREAM, 0));
	time_t ping = E(time(NULL));

	get_ports(argc, argv, &p1, &p2, &p3);
	check_setuid(argc, argv);
	set_port(&sa, p3);
	E(connect(c, (struct sockaddr *) &sa, sizeof(sa)));
	E(setsockopt(c, IPPROTO_TCP, TCP_NODELAY, &t, sizeof(t)));
	fds[PING_IDX] = c;
	E(setsockopt(l, SOL_SOCKET, SO_REUSEADDR, &t, sizeof(t)));
	E(setsockopt(l, SOL_SOCKET, SO_USELOOPBACK, &t, sizeof(t)));
	set_port(&sa, p1);
	E(bind(l, (struct sockaddr *) &sa, sizeof(sa)));
	E(listen(l, 10));
	fds[LISTEN_IDX] = l;
	nfds = 2;
	printf("ready\n");
	fflush(stdout);
	for (;;) {
		fd_set rfdset;
		fd_set efdset;
		int maxfd = 0;
		int i;
		struct timeval tv;

		FD_ZERO(&rfdset);
		FD_ZERO(&efdset);
		for (i = 0; i < nfds; ++i) {
			FD_SET(fds[i], &rfdset);
			FD_SET(fds[i], &efdset);
			if (fds[i] > maxfd) {
				maxfd = fds[i];
			}
		}
		tv.tv_sec = PING_SECS;
		tv.tv_usec = 0;
		E(select(maxfd + 1, &rfdset, NULL, &efdset, &tv));
		if (FD_ISSET(fds[PING_IDX], &rfdset)) {
			if (!E(recv(fds[PING_IDX], buf, sizeof(buf), 0))) {
				fprintf(stderr, "Ping socket closed\n");
				exit(EXIT_SUCCESS);
			}
		}
		if (FD_ISSET(fds[PING_IDX], &efdset)) {
			fprintf(stderr, "Ping socket error\n");
			exit(EXIT_SUCCESS);
		}
		if (FD_ISSET(fds[LISTEN_IDX], &rfdset)) {
			do_accept(l, p2);
		}
		for (i = 2; i < nfds; ++i) {
			/* do_recv() and do_close() may shift the contents of fds */
			if (FD_ISSET(fds[i], &rfdset)) {
				i -= do_recv(i);
			}
			if (FD_ISSET(fds[i], &efdset)) {
				i -= do_close(i);
			}
		}
		if (E(time(NULL)) - ping >= PING_SECS) {
			E(send(fds[PING_IDX], "ping", 4, 0));
			ping = E(time(NULL));
		}
	}
}
